"""
File Name: handWriteDigitsRecognize.py
Author: Chang Zhixuan
Date: 2023-10-30
Description: hand write digits recognize program, including model training and testing.
    MNIST datasets are used.
"""
import pickle

import torch
from torch import optim
import torchvision.transforms.v2
from torchvision import datasets

from utils import *
from Net import *


def load_data(batch_size):
    """
    load data from file.

    :param batch_size: batch size of the data loader
    :return: train_loader, test_loader
    """
    train_set = datasets.MNIST(
        root='./data',
        train=True,
        download=False,
        transform=torchvision.transforms.Compose(
            [
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize((0.1307,), (0.3081,))
            ]
        )
    )
    test_set = datasets.MNIST(
        root='./data',
        train=False,
        download=False,
        transform=torchvision.transforms.Compose(
            [
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize((0.1307,), (0.3081,))
            ]
        )
    )

    # transform function will regularize the input data, and convert the data to tensor
    # data type.
    train_loader = torch.utils.data.DataLoader(
        train_set,
        batch_size=batch_size,
        shuffle=True,
        num_workers=2
    )
    test_loader = torch.utils.data.DataLoader(
        test_set,
        batch_size=batch_size,
        shuffle=False,
        num_workers=2
    )

    # x, y = next(iter(train_loader))
    # print(x.shape, y.shape, x.min(), x.max())
    # plot_image(x, y, 'image sample')

    return train_loader, test_loader


def train(epoch_size, train_loader):
    net = Net()
    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
    train_loss = []
    for epoch in range(epoch_size):
        for batch_idx, (x, y) in enumerate(train_loader):
            # x: [b, 1, 28, 28], y: [512]
            # x: [b, 1, 28, 28] -> [b, 784], out: [b, 10]
            out = net(x.view(x.size(0), 28 * 28))
            y_one_hot = one_hot(y)
            loss = func.mse_loss(out, y_one_hot)

            # vanish gradient
            optimizer.zero_grad()
            # calculate gradient
            loss.backward()
            # w' = w - lr * grad
            optimizer.step()

            train_loss.append(loss.item())

            if batch_idx % 10 == 0:
                print(
                    'epoch index: {}, batch index: {}, loss value: {}'.format(
                        epoch, batch_idx, loss.item()
                    )
                )
    # 绘制损失函数的下降过程
    plot_curve(train_loss)

    return net


def test(net, test_loader):
    correct = 0.
    for x, y in test_loader:
        x = x.view(x.size(0), 28 * 28)
        out = net(x)
        pred = out.argmax(dim=1)
        correct += pred.eq(y).sum().float().item()
    total_num = len(test_loader.dataset)
    acc = correct / total_num
    print('The total accuracy in test dataset is: {}'.format(acc))
    # The total accuracy in test dataset is: 0.9322
    return acc


def main():
    # load_data(batch_size=512)
    batch_size = 512
    train_loader, test_loader = load_data(batch_size)
    net = train(epoch_size=1, train_loader=train_loader)
    acc = test(net, test_loader)
    mdict = {
        'trained_net': net,
        'accuracy': acc
    }
    with open('trained_net2.pybin', 'wb') as fp:
        pickle.dump(mdict, fp)
    return


if __name__ == '__main__':
    main()
